package com.sh.data.engine.domain.shims.hive.util;

import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.hadoop.hive.ql.lib.Node;
import org.apache.hadoop.hive.ql.parse.*;

import java.util.Iterator;
import java.util.List;
import java.util.Set;

/**
 * @author: mengzheng.mhc
 * @date: 2024/6/7 16:24
 */
@Slf4j
public class HiveSqlParserUtil {

    private static final String A = ";";
    private static final String B = "<EOF>";

    public static String getSqlType(String sql) throws ParseException {

        List<ASTNode> nodes = getTree(sql, 0);
        ASTNode child = (ASTNode) nodes.get(0).getChildren().get(0);
        return child.toString();
    }

    public static Set<String> getFunction(String sql) throws ParseException {

        Set<String> list = Sets.newHashSet();

        List<ASTNode> nodes = getTree(sql, HiveParser.TOK_FUNCTION);
        for (ASTNode node : nodes) {
            List<Node> childs = node.getChildren();
            for (Node child : childs) {
                if (((ASTNode) child).getToken().getType() == HiveParser.Identifier) {
                    list.add(child.toString().toLowerCase());
                }
            }
        }
        return list;
    }

    public static List<String> getCreateTableNames(String sql) throws ParseException {

        List<String> list = Lists.newArrayList();

        List<ASTNode> nodes = getTree(sql, HiveParser.TOK_CREATETABLE);
        for (ASTNode node : nodes) {
            List<Node> childs = node.getChildren();
            for (Node child : childs) {
                if (((ASTNode) child).getType() == HiveParser.TOK_TABNAME) {
                    List<Node> tableNames = ((ASTNode) child).getChildren();
                    for (Node tableName : tableNames) {
                        if (((ASTNode) tableName).getType() == HiveParser.Identifier) {
                            list.add(tableName.toString().toLowerCase());
                        }
                    }
                }
            }
        }
        return list;
    }

    public static Set<String> getAllTableNames(String sql) throws ParseException {

        Set<String> list = Sets.newHashSet();

        List<ASTNode> nodes = getTree(sql, HiveParser.TOK_TABNAME);
        for (ASTNode node : nodes) {
            List<Node> childs = node.getChildren();
            for (Node child : childs) {
                if (((ASTNode) child).getType() == HiveParser.Identifier) {
                    list.add(child.toString().toLowerCase());
                }
            }
        }
        return list;
    }

    public static List<ASTNode> getTree(String sql, int tok) throws ParseException {

        List<ASTNode> nodes = Lists.newArrayList();

        ParseDriver pd = new ParseDriver();

        ParseResult result = pd.parse(sql);

        ASTNode tree = result.getTree();

        ergodicTree(tree, nodes, tok);

        return nodes;
    }

    private static void ergodicTree(ASTNode tree, List<ASTNode> nodes, int tok) {

        if ((tree != null && tree.getToken() != null && tree.getToken().getType() == tok)) {
            nodes.add(tree);
        } else if (tok == 0 && tok == tree.getType()) {
            nodes.add(tree);
        }

        if (tree.getChildren() != null) {
            Iterator i = tree.getChildren().iterator();

            while (i.hasNext()) {
                Node node = (Node) i.next();
                if (node instanceof ASTNode) {
                    ergodicTree((ASTNode) node, nodes, tok);
                }
            }
        }
    }

    public static boolean checkSetSql(String sc) {

        if (StringUtils.isBlank(sc)) {
            return false;
        }

        if (sc.toLowerCase().trim().startsWith("set")) {
            return true;
        }

        return false;
    }
}
